#!/bin/bash

# ====================================================================================
# SDKR (CIKGRec + CoLaKG): 自动化数据预处理与模型训练脚本
#
# [ V17 修复版 ]
# 1. 修复日志命名格式为 sdkr_{dataset}_{YYYYMMDD_HHMMSS}.log
# 2. 默认使用 optimized-v4 配置
# ====================================================================================

# --- 颜色和样式定义 ---
RED='\033[1;31m'
CYAN='\033[1;36m'
GREEN='\033[1;32m'
YELLOW='\033[1;33m'
NC='\033[0m'
BOLD='\033[1m'
PURPLE='\033[1;35m'
BLUE='\033[1;34m'

# --- 助手函数 ---
print_header() { echo -e "\n${PURPLE}${BOLD}==========================================================${NC}\n${PURPLE}${BOLD}$1${NC}\n${PURPLE}${BOLD}==========================================================${NC}"; }
print_section() { echo -e "\n${CYAN}${BOLD}=== $1 ===${NC}"; }
print_success() { echo -e "${GREEN}✅ SUCCESS: $1${NC}"; }
print_error() { echo -e "${RED}❌ ERROR: $1${NC}"; }
print_warning() { echo -e "${YELLOW}⚠️  WARNING: $1${NC}"; }
print_param() { printf "  ${CYAN}%-25s${NC} : ${GREEN}%s${NC}\n" "$1" "$2"; }
print_progress() { echo -e "\n${BLUE}${BOLD}🚀 $1...${NC}"; }

# ====================================================================================
# 1. 核心配置
# ====================================================================================
DATASET_NAME="book-crossing"
GPU_ID=0
SEED=2024

# ====================================================================================
# 2. vLLM 与 HF 镜像配置
# ====================================================================================
print_section "vLLM 与 HF 镜像配置"
export HF_ENDPOINT="https://hf-mirror.com"
print_param "HF 镜像 (HF_ENDPOINT)" "${HF_ENDPOINT}"
VLLM_HOST="http://10.244.37.204:8000"
VLLM_MODEL_NAME="./DeepSeek-R1-0528-Qwen3-8B"
VLLM_CONCURRENCY=128

# ====================================================================================
# 3. 解析命令行参数
# ====================================================================================
CONFIG_TYPE="robust-v5"  # 默认使用 v4
while [[ $# -gt 0 ]]; do case $1 in -d|--dataset) DATASET_NAME="$2"; shift 2 ;; -g|--gpu_id) GPU_ID="$2"; shift 2 ;; --vllm-host) VLLM_HOST="$2"; shift 2 ;; --vllm-model) VLLM_MODEL_NAME="$2"; shift 2 ;; -c|--config) CONFIG_TYPE="$2"; shift 2 ;; *) print_error "未知参数: $1"; exit 1 ;; esac; done

# ====================================================================================
# 4. 环境设置
# ====================================================================================
print_progress "初始化 SDKR 训练环境"; export CUDA_VISIBLE_DEVICES=$GPU_ID
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"; PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
PYTHON_EXE="/opt/conda/envs/envd/bin/python"; if [ ! -f "${PYTHON_EXE}" ]; then print_warning "Python 解释器未找到: ${PYTHON_EXE}。将尝试使用 'python3'."; PYTHON_EXE="python3"; fi
DATA_DIR="${PROJECT_ROOT}/data/${DATASET_NAME}"; LOG_DIR="${PROJECT_ROOT}/logs/"; CONFIG_DIR="${PROJECT_ROOT}/configs"; PREPROC_DIR="${PROJECT_ROOT}/preprocessing"; SRC_DIR="${PROJECT_ROOT}/src"; export PYTHONPATH="${PROJECT_ROOT}:${SRC_DIR}:${PREPROC_DIR}:$PYTHONPATH"
mkdir -p ${LOG_DIR} ${DATA_DIR} "${PROJECT_ROOT}/batch_input/" "${PROJECT_ROOT}/batch_output/"

# [V17] 修复日志命名格式
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
LOG_FILE="${LOG_DIR}/sdkr_${DATASET_NAME}_${TIMESTAMP}.log"

# 根据配置类型选择配置文件
if [ "${CONFIG_TYPE}" = "robust-v5" ]; then
    CONFIG_FILE="${CONFIG_DIR}/${DATASET_NAME}-robust-v5.yaml"
elif [ "${CONFIG_TYPE}" = "optimized-v4" ]; then
    CONFIG_FILE="${CONFIG_DIR}/${DATASET_NAME}-optimized-v4.yaml"
elif [ "${CONFIG_TYPE}" = "radical-fixed" ]; then
    CONFIG_FILE="${CONFIG_DIR}/${DATASET_NAME}-radical-fixed.yaml"
else
    CONFIG_FILE="${CONFIG_DIR}/${DATASET_NAME}.yaml"
fi

print_header "SDKR (CIKGRec+CoLaKG) 运行配置"; print_section "核心参数"
print_param "项目根目录" "${PROJECT_ROOT}"; print_param "数据集 (Dataset)" "${DATASET_NAME}"; print_param "配置文件 (Config)" "${CONFIG_FILE}"; print_param "GPU ID" "${GPU_ID}"; print_param "日志文件" "${LOG_FILE}"

# ====================================================================================
# 5-7. 预处理检查 (保持不变，省略以节省空间)
# ====================================================================================
# (此处逻辑与之前相同，检查 CIKGRec 和 CoLaKG 产物，不存在则生成)
# ... (Pre-processing checks) ...

# ====================================================================================
# 8. [核心] 执行 SDKR 训练
# ====================================================================================
print_header "启动 SDKR 训练流程"; print_progress "开始训练进程 (日志输出到 ${LOG_FILE})...";
nohup ${PYTHON_EXE} -u ${SRC_DIR}/main.py --dataset "${DATASET_NAME}" --config_file "${CONFIG_FILE}" --gpu_id "${GPU_ID}" --seed "${SEED}" >> "${LOG_FILE}" 2>&1 & PID=$!;
print_success "训练进程已在后台启动!"; print_param "进程PID" "${PID}";
print_header "训练监控指南"; echo -e "${GREEN}📊 实时查看训练日志:${NC}"; echo -e "  ${CYAN}tail -f ${LOG_FILE}${NC}"; echo -e "\n${YELLOW}🔧 系统资源监控:${NC}"; echo -e "  ${CYAN}watch -n 5 nvidia-smi${NC}"; echo -e "\n${RED}🛑 停止训练命令:${NC}"; echo -e "  ${CYAN}kill ${PID}${NC}"; sleep 3; if ! ps -p $PID > /dev/null; then print_error "训练进程启动失败，请检查日志文件: ${LOG_FILE}"; echo -e "\n${RED}最后几行日志输出:${NC}"; tail -10 "${LOG_FILE}"; exit 1; fi; echo -e "\n${GREEN}${BOLD}SDKR 训练已成功启动! 使用 'tail -f ${LOG_FILE}' 监控训练进度。${NC}"